from ultralytics.data import YOLODataset
from ultralytics.cfg import cfg2dict
from ultralytics.data.augment import Mosaic, MixUp, CopyPaste
import cv2
import random

# 测试数据获取
config = cfg2dict('./coco8/coco8.yaml')
dataset = YOLODataset(data=config, img_path='./coco8/images')

# 将dataset转为labels格式
def dataset2label(idx, dataset, imgsz=False):
    return {"im_file":dataset.get_image_and_label(idx)['im_file'],
            'ori_shape':dataset.get_image_and_label(idx)['ori_shape'] if not imgsz else(640,640),
            "img":dataset.get_image_and_label(idx)['img'] if not imgsz else pad_to_square(dataset.get_image_and_label(idx)['img']),
            "resized_shape":dataset.get_image_and_label(idx)['resized_shape'] if not imgsz else(640,640),
            "cls":dataset.get_image_and_label(idx)['cls'],
            "instances":dataset.get_image_and_label(idx)['instances']}

# 框标签可视化保存
def vis_box(box_list, img, color):
    for box in box_list:
        l,t,r,b = box
        cv2.rectangle(img, (int(l), int(t)), (int(r), int(b)), color, thickness=2, lineType=cv2.LINE_AA)
    return img

# mixup图片大小统一
def pad_to_square(img, imgsz=640):
    height, width = img.shape[:2]
    if height<imgsz:
        h_pad = imgsz-height
    else:
        h_pad = 0
    if width<imgsz:
        w_pad = imgsz-width
    else:
        w_pad = 0
    padded_img = cv2.copyMakeBorder(img, 0, h_pad, 0, w_pad, cv2.BORDER_CONSTANT, value=(114,114,114))
    return padded_img

# Mosaic测试数据labels
labels1 = dataset2label(1, dataset)
labels1['mix_labels'] = [dataset2label(i, dataset) for i in random.choices([0,1,2,3],k=3)]
#print(labels1)

# 初始化Mosaic类
mosaic = Mosaic(dataset, imgsz=640, p=1.0, n=4)
result = mosaic._mosaic4(labels1)
# 图片mosaic变换结果保存
save1dir = './mosaic_test.jpg'
cv2.imwrite(save1dir, result['img'])
# 修正框可视化
box1_list = result['instances']._bboxes.bboxes
vis1dir = './mosaic_vis.jpg'
img1 = cv2.imread(save1dir)
cv2.imwrite(vis1dir, vis_box(box1_list,img1, (0,255,0)))


# MixUp测试数据labels
labels2 = dataset2label(1, dataset, imgsz=True)
labels2['mix_labels'] = [dataset2label(i, dataset, imgsz=True) for i in random.choices([0,2,3],k=1)]  # MixUp只处理两张图片
#print(labels2['instances']._bboxes.bboxes)
# 初始化MixUp类
mixup = MixUp(dataset, p=1.0)
output = mixup._mix_transform(labels2)
# 图片mixup变换结果保存
save2dir = './mixup_test.jpg'
cv2.imwrite(save2dir, output['img'])
